import torch
import torch.nn as nn
import torchvision
from torchvision import transforms
from torch.utils.data import DataLoader

from PIL import Image
import numpy as np
import torch.nn.functional as F

# 数据预处理
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

# 加载CIFAR10数据集
# trainset = torchvision.datasets.CIFAR10('./data', download=True, train=True, transform=transform)
# trainloader = DataLoader(trainset, batch_size=1, shuffle=True, num_workers=0)

# 实例化ResNet50模型 # 自动下载网上的预训练模型
model = torchvision.models.resnet50(pretrained=True)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 20)

# 加载预训练的权重
# checkpoint = torch.load('resnet18-5c106cde.pth')
# model.load_state_dict(checkpoint['model_state_dict'])

# 导入图片并转换为Tensor
image = Image.open("d://Elasticsearch.png")
image_tensor = transform(image).unsqueeze(0)

# 使用GPU（如果可用）
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)
image_tensor.to(device)

# 提取特征
with torch.no_grad():
    output = model(image_tensor)
    logps = F.softmax(output, dim=1)
    probs, classes = torch.max(logps, dim=1)

# 获取特征
feature = model(image_tensor)
feature = feature.cpu().data.numpy()

print(f'Image class: {classes.item()}, Probability: {probs.item()}, Feature: {feature}')
